package me.boot.db.mybatis.datasource;

import cn.hutool.core.util.StrUtil;
import com.baomidou.dynamic.datasource.DynamicRoutingDataSource;
import com.baomidou.dynamic.datasource.creator.DefaultDataSourceCreator;
import com.baomidou.dynamic.datasource.spring.boot.autoconfigure.DataSourceProperty;
import com.baomidou.dynamic.datasource.toolkit.DynamicDataSourceContextHolder;
import me.boot.db.mybatis.propertites.BasicDataSourceProperties;
import org.springframework.beans.BeanUtils;

import javax.sql.DataSource;
import java.util.Map;
import java.util.function.Function;

/**
 * 多数据源管理服务
 */
public interface DynamicDataSourceService {


    /**
     * 按约定规则，获取租户的独立数据源名称
     *
     * @param tenantId 租户ID
     * @return 数据源名称
     */
    static String getTenantDsName(String tenantId) {
        return StrUtil.format("tenant_{}", tenantId);
    }


    /**
     * 切换到租户的独立数据源，
     *
     * @param tenantId 租户ID
     * @param function 在切换数据源后，执行调用方的业务函数，入参是数据源名称
     * @param <R>      业务函数返回值的类型
     * @return 业务函数返回值，即function的执行结果
     * @throws Throwable
     */
    default <R> R holdTenant(String tenantId, Function<String, R> function) throws Throwable {
        String dsName = DynamicDataSourceService.getTenantDsName(tenantId);
        if (contains(dsName)) {
            try {
                DynamicDataSourceContextHolder.push(dsName);
                return function.apply(dsName);
            } finally {
                DynamicDataSourceContextHolder.poll();
            }
        } else {
            throw new Throwable(StrUtil.format("租户ID={} 的数据源={} 未加载", tenantId, dsName));
        }
    }


    /**
     * 是否存在指定数据源
     *
     * @param dsName 数据源名称
     * @return 存在标记
     */
    default boolean contains(String dsName) {
        return getAll().containsKey(dsName);
    }

    /**
     * 获取当前数据源名称
     *
     * @return 当前数据源名称
     */
    default String getCurrentDsName() {
        return DynamicDataSourceContextHolder.peek();
    }

    /**
     * 获取当前数据源
     *
     * @return 当前数据源
     */
    default DataSource getCurrentDs() {
        return getDs(getCurrentDsName());
    }


    /**
     * 获取数据源
     *
     * @param key 数据源名称
     * @return 数据源
     */
    default DataSource getDs(String key) {
        return getAll().get(key);
    }

    /**
     * 动态添加数据源
     *
     * @param property 数据源属性简易描述
     */
    default void addDs(String dsName, BasicDataSourceProperties property) {
        DataSourceProperty dataSourceProperty = new DataSourceProperty();
        BeanUtils.copyProperties(property, dataSourceProperty);
        DataSource dataSource = getDsCreator().createDataSource(dataSourceProperty);
        getDynamicRoutingDs().addDataSource(dsName, dataSource);
    }

    /**
     * 动态移除数据源
     *
     * @param key
     */
    default void removeDs(String key) {
        getDynamicRoutingDs().removeDataSource(key);
    }

    /**
     * 获取所有数据源
     * key   -- 数据源名称
     * value -- 数据源
     *
     * @return 所有数据源
     */
    default Map<String, DataSource> getAll() {
        return getDynamicRoutingDs().getDataSources();
    }


    /**
     * 获取动态路由数据源
     *
     * @return 动态路由数据源
     */
    DynamicRoutingDataSource getDynamicRoutingDs();

    /**
     * 获取默认的数据源创建器
     *
     * @return 数据源创建器
     */
    DefaultDataSourceCreator getDsCreator();
}
